#include "search/search_server.hpp"
#include "reactor/event_loop.hpp"
#include "reactor/tcp_connection.hpp"
#include "reactor/buffer.hpp"
#include <iostream>
#include <nlohmann/json.hpp> 
#include "external/utfcpp/utf8.h"

using json = nlohmann::json;

// 函数：清理字符串，确保其为合法的 UTF-8 编码
std::string sanitize_utf8(const std::string& input) {
    std::string clean_str;
    clean_str.reserve(input.length());
    
    auto it = input.begin();
    auto end = input.end();
    
    while (it != end) {
        auto next_it = utf8::find_invalid(it, end);
        clean_str.append(it, next_it);
        it = next_it;
        if (it != end) {
            // 跳过这个无效字节
            ++it;
        }
    }
    return clean_str;
}

SearchServer::SearchServer(const std::string& ip, uint16_t port, int threadNum)
    : _mainLoop(new reactor::EventLoop()),
      _server(_mainLoop.get(), ip, port) 
{
    _server.setThreadNum(threadNum);
    _server.setConnectionCallback(
        std::bind(&SearchServer::onConnection, this, std::placeholders::_1));
    _server.setMessageCallback(
        std::bind(&SearchServer::onMessage, this, std::placeholders::_1, std::placeholders::_2));
}

SearchServer::~SearchServer() {}

bool SearchServer::init(const std::string& cnDict, const std::string& cnIndex,
                        const std::string& enDict, const std::string& enIndex,
                        const std::string& invertedIndex, const std::string& pages,
                        const std::string& offsets) {
    std::cout << "--- 开始加载关键字推荐数据 ---" << std::endl;
    if (!_keywordRecommend.init(cnDict, cnIndex, enDict, enIndex)) {
        std::cerr << "关键字推荐模块数据加载失败！" << std::endl;
        return false;
    }
    std::cout << "--- 关键字推荐数据加载完成 ---\n" << std::endl;

    std::cout << "--- 开始加载网页搜索数据 ---" << std::endl;
    if (!_pageSearch.init(invertedIndex, pages, offsets)) {
        std::cerr << "网页搜索模块数据加载失败！" << std::endl;
        return false;
    }
    std::cout << "--- 网页搜索数据加载完成 ---" << std::endl;
    
    return true;
}

void SearchServer::start() {
    std::cout << "\n服务器启动于 " << _server.ip() << ":" << _server.port() << std::endl;
    _server.start();
    _mainLoop->loop();
}

void SearchServer::onConnection(const std::shared_ptr<reactor::TcpConnection>& conn) {
    if (conn->connected()) {
        std::cout << "新连接: " << conn->sockfd() << std::endl;
        std::lock_guard<std::mutex> lock(_parsersMutex);
        _parsers[conn->sockfd()] = protocol::MessageParser();
    } else {
        std::cout << "连接关闭: " << conn->sockfd() << std::endl;
        std::lock_guard<std::mutex> lock(_parsersMutex);
        _parsers.erase(conn->sockfd());
    }
}

void SearchServer::onMessage(const std::shared_ptr<reactor::TcpConnection>& conn, reactor::Buffer* buf) {
    protocol::MessageParser* parser = nullptr;
    {
        std::lock_guard<std::mutex> lock(_parsersMutex);
        if (_parsers.find(conn->sockfd()) == _parsers.end()) {
            return;
        }
        parser = &_parsers.at(conn->sockfd());
    }
    
    size_t processedBytes = 0;
    std::string receivedData = buf->retrieveAllAsString();
    std::vector<protocol::Message> messages = parser->parse(receivedData.c_str(), receivedData.size(), processedBytes);
    
    for (const auto& msg : messages) {
        std::string responseContent;
        protocol::MsgType responseType;

        if (msg.tag == static_cast<int>(protocol::MsgType::KEYWORD_RECOMMEND_REQ)) {
            responseContent = handleKeywordRecommend(msg.value);
            responseType = protocol::MsgType::KEYWORD_RECOMMEND_RESP;
        } else if (msg.tag == static_cast<int>(protocol::MsgType::PAGE_SEARCH_REQ)) {
            responseContent = handlePageSearch(msg.value);
            responseType = protocol::MsgType::PAGE_SEARCH_RESP;
        } else {
            json error_json = {{"error", "无效的请求类型"}};
            responseContent = error_json.dump();
            responseType = protocol::MsgType::ERROR_RESP;
        }
        
        protocol::Message response;
        response.tag = static_cast<int32_t>(responseType);
        response.value = responseContent;
        response.length = responseContent.size();
        conn->send(response.serialize());
    }
}

std::string SearchServer::handleKeywordRecommend(const std::string& query) {
    auto results = _keywordRecommend.recommend(query);
    json j_results = json::array();
    for (const auto& [word, freq] : results) {
        j_results.push_back({{"word", word}, {"frequency", freq}});
    }
    return j_results.dump();
}

std::string SearchServer::handlePageSearch(const std::string& query) {
    auto results = _pageSearch.search(query);
    json j_results = json::array();
    for (const auto& r : results) {
        j_results.push_back({
            {"title", sanitize_utf8(r.title)},       // **调用清理函数**
            {"link", sanitize_utf8(r.link)},         // link也清理一下，更安全
            {"abstract", sanitize_utf8(r.abstract)}, // **调用清理函数**
            {"score", r.score}
        });
    }
    return j_results.dump();
}
